[Relax][Frontend][TFLite] Support dynamic RANGE scalar bounds#19867
[Relax][Frontend][TFLite] Support dynamic RANGE scalar bounds#19867Aharrypotter wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for TFLite RANGE operators with dynamic (runtime) scalar bounds in the Relax frontend, handling both integer and float data types by computing the element count in-graph and lifting it to a symbolic dimension. It also replaces the previous unsupported-feature test with comprehensive parameterized tests. Feedback on the changes points out a critical typo where tirx is used instead of tir in the _scalar_tensor_to_dim helper, which will cause a NameError at runtime.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Aharrypotter
force-pushed
the
tflite-range-dynamic-19412
branch
from
a359f03 to
c9bd9ca
Compare
Previously convert_range raised OpNotImplemented when start/limit/delta were runtime (non-constant) scalar tensors, so any TFLite model that computes RANGE bounds at runtime failed to import. This was one of the "partial implementation" items tracked in apache#19412. relax.op.arange only takes compile-time PrimExpr bounds, and its struct-info length formula (InferTypeArange) has no negative-step branch, so feeding symbolic bounds straight in would mis-declare descending ranges. Instead, compute the element count in-graph and lift it to one symbolic output dimension via relax.op.tensor_to_shape + match_cast (the bridge already used by _get_shape_expr_from_tensor), so the declared and runtime lengths match by construction. Values are rebuilt as arange(0, count) * delta + start. One unified path covers both dtypes: - int: count = -floor_divide(start - limit, delta), exact and sign-agnostic (no float-precision loss), equal to ceil((limit - start) / delta); - float: count = ceil((limit - start) / delta). No new Relax op is needed. Replace the "not supported" test with a compile-and-run test covering ascending/descending integer and float dynamic bounds.
Aharrypotter
force-pushed
the
tflite-range-dynamic-19412
branch
from
c9bd9ca to
8dfc70e
Compare
1 participant
Summary
This PR adds Relax TFLite frontend support for dynamic (runtime) scalar bounds
in the
RANGEoperator, addressing theRANGE"fix partial implementations"item from #19412 section C.
convert_rangepreviously lowered only constantstart,limit, anddeltatorelax.op.arangeand raisedOpNotImplementedfor runtime scalarbounds (the guard added in #19401). Models that compute RANGE bounds at runtime
could therefore not be imported. This PR makes the dynamic path work for both
integer and float bounds, ascending or descending, without adding a new Relax
op. The change is limited to the
RANGEconverter and its test.#19813 added a batch of missing TFLite operator mappings but did not touch this
partial-implementation item; this PR closes it.
Design
Dynamic scalar bounds via count-lift
relax.op.arangeonly accepts compile-timePrimExprbounds. The frontendalready has a runtime-scalar -> symbolic-dimension bridge
(
relax.op.tensor_to_shape+match_cast, as used by_get_shape_expr_from_tensor), so no new op is needed.Rather than feed symbolic bounds straight into
arange, the converter computesthe element count in-graph and lifts that single value to one symbolic
output dimension
L, then rebuilds the values asarange(0, L) * delta + start.Lifting the count (instead of the bounds) keeps the declared and runtime output
lengths equal by construction:
arange's struct-info length formula(
InferTypeArange) has no negative-step branch, so feeding symbolic boundsdirectly would mis-declare descending ranges relative to the TOPI runtime
length.
The count is
max(0, ceil((limit - start) / delta)), computed per dtype:-floor_divide(start - limit, delta)— exact, sign-agnostic, andfree of float-precision loss; equal to
ceil((limit - start) / delta).ceil((limit - start) / delta).Constant (all-bounds-constant) RANGE keeps the existing direct-
arangepathunchanged.
Operator Support
RANGEstart,limit,deltarelax.op.arange(constant bounds); count-lift +arange(0, L) * delta + start(dynamic bounds)Tests
The dynamic test compiles the imported module and runs it on the Relax VM,
comparing the output against
numpy.arange. The constant-bound structural testis unchanged.
test_rangetest_range_dynamic_scalar_inputsLocal validation:
Result:
References
RANGE)RANGEdynamic-scalar guard and its testconvert_range